# 用于测量粉末颗粒的尺寸和形状

from skimage import io, filters, measure, segmentation, morphology, color
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import joblib
from pathlib2 import Path
import warnings

# 忽略警告
warnings.filterwarnings('ignore')

# 测量参数名称
region_props = ['area', 'convex_area', 'bbox_area', 'filled_area', 'major_axis_length', 'minor_axis_length', 'moments_hu',
                'eccentricity', 'equivalent_diameter', 'extent', 'feret_diameter_max', 'perimeter', 'solidity', 'perimeter_crofton',
                'orientation', 'inertia_tensor_eigvals', 'bbox', 'centroid', 'image', 'label']

stand = joblib.load(r'C:\Users\admin\Desktop\qtwork\粉\stand.pkl')
model = joblib.load(r'C:\Users\admin\Desktop\qtwork\粉\model.pkl')

class PowderMeasurement:
    def __init__(self, file_path, scale=1, bins=10):
        self.file_path = file_path
        self.scale = scale
        self.bins = bins
    
    def run(self):
        cwd = Path(self.file_path)
        image_files = cwd.glob('*.jpg')
        df_data = pd.DataFrame()
        i = 1
        for image_file in image_files:
            if 'FG' not in str(image_file):
                df = self.process_image(image_file)
                if df.shape[0] > 2:
                    df_data = pd.concat([df_data, df])

            yield i
            i += 1
        data_name = self.file_path + '\\data.csv'
        df_data.to_csv(data_name, index=False)
        
        # 统计结果
        if df_data.shape[0] > 0:
            df_stat = self.statistics(df_data, self.scale,self.bins)
            stat_name = self.file_path + '\\stat.csv'
            df_stat.to_csv(stat_name, index=False)
            self.result = df_data
        else:
            n = np.random.randint(1, 100, size=(5,25))
            self.result = pd.DataFrame(n)
        
        yield 100
    
    # 处理 图像
    def process_image(image_path):
        # 读取图像
        image = io.imread(image_path)
        
        # 转换为灰度图像
        if image.shape[2] == 3:
            gray = color.rgb2gray(image)
        else:
            gray = image
        
        # 二值化
        bw = gray < (80/255)
        gray = color.rgb2gray(image)
        
        # 二值化
        bw = gray < (80/255)
        bw = morphology.remove_small_objects(bw, min_size=25)
        bw = segmentation.clear_border(bw)
        image_new = bw.copy()

        bw_label = measure.label(bw)
        regions = measure.regionprops_table(bw_label, properties=region_props)
        df = pd.DataFrame(regions)

        # 机器学习筛选
        X = df.loc[:, 'area':'inertia_tensor_eigvals-1'].to_numpy()
        X_trans = stand.transform(X)
        y_pred = model.predict(X_trans)
        df['pred'] = y_pred

        image_new_name = image_path.split('.')[0] + '_new.png'
        io.imsave(image_new_name, image_new)

        # sobel 筛选
        sobel = filters.sobel(gray)
        for index, row in df.iterrows():
            xmin, ymin, xmax, ymax = row['bbox-0':'bbox-3']
            region = sobel[ymin:ymax, xmin:xmax].max()
            df.loc[index, 'sobel'] = region
            if (row['pred'] == 1) and (region < 10):
                image_new[bw_label == index + 1] = False
                df.loc[index, 'remove'] = True
            else:
                df.loc[index, 'remove'] = False

        # 二值化图像转换并保存
        image_new = image_new.astype(np.uint8)
        io.imsave(image_new_name, image_new)
        
        return df

    # 统计结果和分级
    def statistics(self, df, scale, bin_step:int):
        area = df['area'].to_numpy()
        dia = np.sqrt(4 * area / np.pi) * scale
        if bin_step == 10:
            dia_bin = np.array([0, 10, 20, 30, 40, 50, 60, 70, 75, 80, 90, 100, 110, 120])
        else:
            dia_low = np.min(dia) // bin_step * bin_step
            dia_high = np.max(dia) // bin_step * bin_step + bin_step
            dia_bin = np.arange(dia_low, dia_high, bin_step)
            
        hist, _ = np.histogram(dia, bins=dia_bin)
        hist = np.round(hist / len(dia), 2)
        df = pd.DataFrame({'dia': dia_bin, 'count': hist})
        # 计算统计值
        # 超过120的百分比，保留 两位小数
        over_120 = np.round(np.sum(dia >= 120) / len(dia) *100, 2)
        df = df.append({'dia': 120, 'count': over_120}, ignore_index=True)
        # 0-75的百分比，保留 两位小数
        under_75 = np.round(np.sum(dia < 75) / len(dia) *100, 2)
        df = df.append({'dia': 75, 'count': under_75}, ignore_index=True)
        # 平均值
        mean = np.round(np.mean(dia), 2)
        df = df.append({'dia': 'mean', 'count': mean}, ignore_index=True)
        # 最大值
        max_dia = np.round(np.max(dia), 2)
        df = df.append({'dia': 'max', 'count': max_dia}, ignore_index=True)
        # 最小值
        min_dia = np.round(np.min(dia), 2)
        df = df.append({'dia': 'min', 'count': min_dia}, ignore_index=True)

        return df

    def read_images(self,file_path, scale, bin_step):
        # 读取文件夹中的所有图像
        cwd = Path(file_path)
        image_files = cwd.rglob('*.jpg')
        df_data = pd.DataFrame()
        for image_file in image_files:
            if 'FG' not in str(image_file):
                df = self.process_image(image_file)
                if df.shape[0] > 2:
                    df_data = pd.concat([df_data, df])
        data_name = file_path + '\\data.csv'
        df_data.to_csv(data_name, index=False)
        
        # 统计结果
        df_stat = self.statistics(df_data, scale, bin_step)
        stat_name = file_path + '\\stat.csv'
        df_stat.to_csv(stat_name, index=False)
    